
import os
import SimpleITK as sitk
from PIL import Image
import numpy as np
import scipy.misc
import cv2
# from sklearn import
slice_number = 24
slice_size = 256

def img_center_crop(image, crop_size):
    assert len(image.shape) == 3, 'invalid image size in sliding window'

    z_start, x_start, y_start = 0, 0, 0
    img_z, img_x, img_y = image.shape[0], image.shape[1], image.shape[2]
    crop_z, crop_x, crop_y = crop_size[0], crop_size[1], crop_size[2]
    # if img_x>set_size and img_y>set_size
    # x or y 一个比crop大 或者两个都大
    if img_x > crop_x and img_y > crop_y:
        starting = [int((crop_z - img_z) / 2), int((img_x - crop_x) / 2), int((img_y - crop_y) / 2)]
        z_start, x_start, y_start = starting[0], starting[1], starting[2]

    elif img_x > crop_x and img_y <= crop_y:
        starting = [int((crop_z - img_z) / 2), int((img_x - crop_x) / 2), int((crop_y - img_y) / 2)]
        z_start, x_start, y_start = starting[0], starting[1], 0
        y_up = int((crop_y - img_y) / 2)
        y_down = crop_y - img_y - y_up
        image = np.pad(image, ((0, 0), (0, 0), (y_up, y_down)), mode='constant')

    elif img_x <= crop_x and img_y > crop_y:
        starting = [int((crop_z - img_z) / 2), int((crop_x - img_x) / 2), int((img_y - crop_y) / 2)]
        z_start, x_start, y_start = starting[0], 0, starting[2]
        x_up = int((crop_x - img_x) / 2)
        x_down = crop_x - img_x - x_up
        image = np.pad(image, ((0, 0), (x_up, x_down), (0, 0)), mode='constant')

    img_crop = image[z_start: z_start + crop_size[0], x_start:x_start + crop_size[1],
               y_start: y_start + crop_size[2]]

    return img_crop


def resample_image(itk_image, out_spacing=(1., 1., 1.), is_label=False):
    original_spacing = itk_image.GetSpacing()
    original_size = itk_image.GetSize()
    out_size = [int(np.round(original_size[0] * (original_spacing[0] / out_spacing[0]))),
                int(np.round(original_size[1] * (original_spacing[1] / out_spacing[1]))),
                int(np.round(original_size[2] * (original_spacing[2] / out_spacing[2])))]

    resample = sitk.ResampleImageFilter()
    resample.SetOutputSpacing(out_spacing)
    resample.SetSize(out_size)
    resample.SetOutputDirection(itk_image.GetDirection())
    resample.SetOutputOrigin(itk_image.GetOrigin())
    resample.SetTransform(sitk.Transform())
    resample.SetDefaultPixelValue(itk_image.GetPixelIDValue())

    if is_label:
        resample.SetInterpolator(sitk.sitkNearestNeighbor)
    else:
        resample.SetInterpolator(sitk.sitkBSpline)

    return resample.Execute(itk_image)

def padding_image_array_size(image_array, out_size):
    img_z, img_x, img_y = image_array.shape[0], image_array.shape[1], image_array.shape[2]
    out_z, out_x, out_y = out_size[0], out_size[1], out_size[2]
    if out_z > img_z:
        z_up = int((out_z - img_z) / 2)
        z_down = out_z - img_z - z_up

        if out_x >= img_x and out_y >= img_y:  # 三个维度都是padding
            x_up = int((out_x - img_x) / 2)
            x_down = out_x - img_x - x_up
            y_up = int((out_y - img_y) / 2)
            y_down = out_y - img_y - y_up
            new_volume = np.pad(image_array, ((z_up, z_down), (x_up, x_down), (y_up, y_down)), mode='constant')
        else:
            new_volume = np.pad(image_array, (z_up, z_down), mode='constant')
            new_volume = img_center_crop(new_volume, (slice_number, slice_size, slice_size))
    else:
        # 把z轴crop为slice number
        z_start = int((img_z - out_z) / 2)
        image_array = image_array[z_start: z_start + out_size[0], :, :]
        if out_x >= img_x and out_y >= img_y:  # 俩个维度都是padding
            x_up = int((out_x - img_x) / 2)
            x_down = out_x - img_x - x_up
            y_up = int((out_y - img_y) / 2)
            y_down = out_y - img_y - y_up
            new_volume = np.pad(image_array,((0,0),(x_up, x_down), (y_up, y_down)), mode='constant')
        else:
            new_volume = img_center_crop(image_array, (24, 256, 256))

    return new_volume

def resample_standard_list(itk_list):
    itk_list = [resample_standard(itk) for itk in itk_list]
    return itk_list

def resample_standard(itk_image):
    resampled_image = resample_image(itk_image, out_spacing=(1., 1., 1.),is_label=False)  # itk_image.GetSize (x,y,z)
    resampled_image = sitk.GetArrayFromImage(resampled_image)  # GetArrayFromImage (z,x,y)
    image_resample = padding_image_array_size(resampled_image, out_size=(24, 256, 256))
    itk_image_resample = sitk.GetImageFromArray(image_resample)

    return itk_image_resample

    # reader = sitk.ImageSeriesReader()
    # series_IDs = reader.GetGDCMSeriesIDs(patient_path)
    # file_reader = sitk.ImageFileReader()
    # # dcm_series = reader.GetGDCMSeriesFileNames(patient_path)
    # for series in series_IDs:
    #     series_file_names = reader.GetGDCMSeriesFileNames(patient_path, series)
    #     # 根据一个单张的dcm文件，读取这个series的metedata，即可以获取这个序列的描述符
    #     file_reader.SetFileName(series_file_names[0])
    #     file_reader.ReadImageInformation()
    #     series_description = file_reader.GetMetaData("0008|103e")
    #     if 't1' in series_description and 'tra' in series_description:
    #         itk_image = sitk.ReadImage(series_file_names)
    #         resampled_image = resample_image(itk_image, out_spacing=(1., 1., 1.),is_label=False)  # itk_image.GetSize (x,y,z)
    #         resampled_image = sitk.GetArrayFromImage(resampled_image,is_label=False)  # GetArrayFromImage (z,x,y)
    #         image_resample = padding_image_array_size(resampled_image, out_size=(24, 256, 256))
    #         itk_image_resample = sitk.GetImageFromArray(image_resample)
    #         sitk.WriteImage(itk_image_resample, new_patient_path+'/t1.nii.gz')

    #     if 't2' in series_description and 'tra' in series_description and 'dark-fluid' in series_description:
    #         itk_image = sitk.ReadImage(series_file_names)
    #         resampled_image = resample_image(itk_image, out_spacing=(1., 1., 1.),is_label=False)  # itk_image.GetSize (x,y,z)
    #         resampled_image = sitk.GetArrayFromImage(resampled_image)  # GetArrayFromImage (z,x,y)
    #         image_resample = padding_image_array_size(resampled_image, out_size=(24, 256, 256))
    #         itk_image_resample = sitk.GetImageFromArray(image_resample)
    #         sitk.WriteImage(itk_image_resample, new_patient_path + '/t2.nii.gz')

